﻿using System;
using RayDen.Library.Components.Surface;
using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;
using RayDen.RayEngine.Core.Interface;
using RayDen.RayEngine.Core.Types;
using RayDen.RayEngine.Core.Types.Volume;

namespace RayDen.RayEngine.Engines.PathTracer.PathSamplers
{
    public class PathTracerSampler : PathTracerPathBase<PathTracerPathState>
    {

        protected VolumeComputation volumeComp;
        protected RayHit eyeHit;

        public override void InitPath(IPathProcessor buffer)
        {
            base.InitPath(buffer);
            if (scene.VolumeIntegrator != null)
            {

                if (this.volumeComp == null)
                {
                    this.volumeComp = scene.GetVolumeComputation();
                }
                else
                {
                    this.volumeComp.Reset();
                }
            }
            this.PathState = PathTracerPathState.EyeVertex;
        }


        protected override int RaysOccupied()
        {
            switch (PathState)
            {

                case PathTracerPathState.EyeVolume:
                    return 1 + (volumeComp == null ? 0 : volumeComp.GetRayCount());
                case PathTracerPathState.EyeVertex:
                    return 1;
                case PathTracerPathState.NextVertex:
                    return 1 + tracedShadowRayCount;
                case PathTracerPathState.ShadowRaysOnly:
                    return tracedShadowRayCount;
                default:
                    throw new ArgumentException("Invalid ray state");
            }
        }

        protected override void EnqueueRays(RayBuffer rb)
        {
            if (PathState == PathTracerPathState.EyeVolume)
            {
                volumeComp.AddRays(rb);
            }
            if (PathState != PathTracerPathState.ShadowRaysOnly)
                RayIndex = rb.AddRay(PathRay);
            if (PathState == PathTracerPathState.NextVertex || PathState == PathTracerPathState.ShadowRaysOnly)
            {
                for (int i = 0; i < tracedShadowRayCount; ++i)
                    secRays[i].currentShadowRayIndex = rb.AddRay(secRays[i].shadowRay);
            }
        }


        protected override bool SampleSecondaryRays(ref Point hitPoint, ref Vector wo)
        {
            tracedShadowRayCount = 0;

            float lightStrategyPdf = LightSampler.StrategyPdf;
            //scene.ShadowRaysPerSample/(float)scene.Lights.Length;
            RgbSpectrum lightTroughtput = Throughput * hitInfo.Color;
            LightSampler.EvaluateShadow(ref hitPoint, ref hitInfo.Normal, Sample.GetLazyValue(),
                                        Sample.GetLazyValue(), Sample.GetLazyValue(), ref ls);
            for (int index = 0; index < ls.Length; index++)
            {

                if (ls[index].Pdf <= 0f)
                    continue;
                secRays[tracedShadowRayCount].color = (RgbSpectrum)(ls[index].Spectrum);
                secRays[tracedShadowRayCount].pdf = ls[index].Pdf;
                secRays[tracedShadowRayCount].shadowRay = ls[index].LightRay;
                Vector lwi = secRays[tracedShadowRayCount].shadowRay.Dir;
                RgbSpectrum fs;
                hitInfo.MMaterial.f(ref secRays[tracedShadowRayCount].shadowRay.Dir, ref wo, ref hitInfo.ShadingNormal, ref Throughput, out fs, types: BrdfType.Diffuse);
                secRays[tracedShadowRayCount].color *= lightTroughtput *
                                                            Vector.AbsDot(ref hitInfo.Normal, ref lwi) *
                                                            fs;
                if (!secRays[tracedShadowRayCount].color.IsBlack())
                {
                    secRays[tracedShadowRayCount].pdf /= lightStrategyPdf;
                    tracedShadowRayCount++;
                }
            }
            /*
            float lightStrategyPdf = scene.ShadowRayCount / (float)scene.Lights.Length;
            RgbSpectrum lightTroughtput = Throughput * hitInfo.Color;
            for (int i = 0; i < scene.ShadowRaysPerSample; ++i)
            {
                int currentLightIndex = scene.SampleLights(Sample.GetLazyValue());
                var light = scene.Lights[currentLightIndex];

                var ls = new LightSample();
                light.EvaluateShadow(ref hitPoint, ref hitInfo.ShadingNormal, Sample.GetLazyValue(),
                                     Sample.GetLazyValue(), Sample.GetLazyValue(), out ls);
                if (ls.Pdf <= 0f)
                {
                    continue;
                }

                secRays[tracedShadowRayCount].color = new RgbSpectrum(ls.Spectrum);
                secRays[tracedShadowRayCount].pdf = ls.Pdf;
                secRays[tracedShadowRayCount].shadowRay = ls.LightRay;

                RgbSpectrum fs;

                hitInfo.MMaterial.f(
                    ref secRays[tracedShadowRayCount].shadowRay.Dir,
                    ref wo, ref hitInfo.ShadingNormal, out fs);
                Vector lwi = secRays[tracedShadowRayCount].shadowRay.Dir;
                // *scene.Transmittance !!!TODO
                secRays[tracedShadowRayCount].color *= lightTroughtput *
                                                       Vector.AbsDot(ref hitInfo.ShadingNormal, ref lwi) *
                                                       fs;
                if (!secRays[tracedShadowRayCount].color.IsBlack())
                {
                    secRays[tracedShadowRayCount].pdf *= lightStrategyPdf;
                    tracedShadowRayCount++;
                }
            }
             */ 
            return true;
        }

        protected override bool OnSecondaryRays(RayBuffer rayBuffer, SampleBuffer sb)
        {
            if (((PathState == PathTracerPathState.ShadowRaysOnly)) || ((PathState == PathTracerPathState.NextVertex) && (tracedShadowRayCount > 0)))
            {
                for (int i = 0; i < tracedShadowRayCount; ++i)
                {
                    RayHit shadowRayHit = rayBuffer.rayHits[secRays[i].currentShadowRayIndex];
                    RgbSpectrum attenuation;
                    bool continueTrace;
                    if (this.ShadowRayTest(ref shadowRayHit, ref secRays[i].shadowRay, out attenuation, out continueTrace))
                    {
                        //                            Radiance.MADD()
                        Radiance += ((secRays[i].color) / secRays[i].pdf);
                      //  pathWeight *= secRays[i].pdf;
                    }
                }

            }

            switch (PathState)
            {
                case PathTracerPathState.EyeVertex:
                    float t, t1;
                    if (scene.VolumeIntegrator != null && scene.VolumeIntegrator.GetBounds().Intersect(PathRay, out t, out t1))
                    {
                        var rayHit = rayBuffer.rayHits[RayIndex];

                        // Use Russian Roulette to check if I have to do participating media computation or not
                        if (Sample.GetLazyValue() <= scene.VolumeIntegrator.GetRRProbability())
                        {
                            var  volumeRay = new RayData(PathRay.Org, PathRay.Dir, 0f,
                                                            rayHit.Miss() ? 1e10f : rayHit.Distance);
                            scene.VolumeIntegrator.GenerateLiRays(scene, Sample, ref volumeRay, volumeComp);
                            Radiance += volumeComp.GetEmittedLight();

                            if (volumeComp.GetRayCount() > 0)
                            {
                                // Do the EYE_VERTEX_VOLUME_STEP
                                PathState = PathTracerPathState.EyeVolume;
                                eyeHit = (rayBuffer.rayHits[RayIndex]);
                                return true;
                            }
                        }
                    }

                    break;
                    //case PathTracerPathState.NextVertex:rayHit = rayBuffer.rayHits[RayIndex];
                case PathTracerPathState.EyeVolume:
                    Radiance += Throughput * volumeComp.CollectResults(rayBuffer) /
                                scene.VolumeIntegrator.GetRRProbability();
                    //rayHit = eyeHit;
                    break;
            }

            if (PathState == PathTracerPathState.ShadowRaysOnly)
            {
                Splat(sb);
                return true;
            }

            return false;
        }

        protected override bool OnBsdfSampling(ref RgbSpectrum f, ref float fPdf)
        {

            if ((fPdf <= 0.0f) || f.IsBlack())
            {
                if (tracedShadowRayCount > 0)
                    PathState = PathTracerPathState.ShadowRaysOnly;
                else
                {
                    return false;
                }
            }

            PathState = PathTracerPathState.NextVertex;
            pathWeight *= fPdf;
            Throughput *= f / fPdf;
            return true;
        }


        protected override bool ExtendPath(ref RgbSpectrum f, float p)
        {

            if (depth > scene.MaxPathDepth)
            {
                float prob = Math.Max(Throughput.Filter(), scene.RussianRuletteImportanceCap);
                if (prob >= Sample.GetLazyValue())
                {
                    Throughput /= prob;
                    pathWeight *= prob;
                }
                else
                {
                    if (tracedShadowRayCount > 0)
                        PathState = PathTracerPathState.ShadowRaysOnly;
                    else
                    {
                        return false;
                    }
                }
            }
            return true;
        }

        protected override void OnEnvironmentHit(ref Vector dir, RayBuffer rayBuffer, SampleBuffer consumer)
        {
            if (specularBounce)
            {
                Radiance += scene.SampleEnvironment(ref dir) * Throughput;
            }
        }
    }
}